package hnx.ca.util;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.OutputStream;
import java.security.PrivateKey;
import java.security.Provider;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;

import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.KeyValue;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.w3c.dom.Document;

/**
 * This is a simple example of generating an Enveloped XML 
 * Signature using the JSR 105 API. The resulting signature will look 
 * like (key and signature values will be different):
 *
 * <pre><code>
 *<Envelope xmlns="urn:envelope">
 * <Signature xmlns="http://www.w3.org/2000/09/xmldsig#">
 *   <SignedInfo>
 *     <CanonicalizationMethod Algorithm="http://www.w3.org/TR/2001/REC-xml-c14n
-20010315"/>
 *     <SignatureMethod Algorithm="http://www.w3.org/2000/09/xmldsig#dsa-sha1"/>
 *     <Reference URI="">
 *       <Transforms>
 *         <Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
 *       </Transforms>
 *       <DigestMethod Algorithm="http://www.w3.org/2000/09/xmldsig#sha1"/>
 *       <DigestValue>K8M/lPbKnuMDsO0Uzuj75lQtzQI=<DigestValue>
 *     </Reference>
 *   </SignedInfo>
 *   <SignatureValue>
 *     DpEylhQoiUKBoKWmYfajXO7LZxiDYgVtUtCNyTgwZgoChzorA2nhkQ==
 *   </SignatureValue>
 *   <KeyInfo>
 *     <KeyValue>
 *       <DSAKeyValue>
 *         <P>
 *           rFto8uPQM6y34FLPmDh40BLJ1rVrC8VeRquuhPZ6jYNFkQuwxnu/wCvIAMhukPBL
 *           FET8bJf/b2ef+oqxZajEb+88zlZoyG8g/wMfDBHTxz+CnowLahnCCTYBp5kt7G8q
 *           UobJuvjylwj1st7V9Lsu03iXMXtbiriUjFa5gURasN8=
 *         </P>
 *         <Q>
 *           kEjAFpCe4lcUOdwphpzf+tBaUds=
 *         </Q>
 *         <G>
 *           oe14R2OtyKx+s+60O5BRNMOYpIg2TU/f15N3bsDErKOWtKXeNK9FS7dWStreDxo2
 *           SSgOonqAd4FuJ/4uva7GgNL4ULIqY7E+mW5iwJ7n/WTELh98mEocsLXkNh24HcH4
 *           BZfSCTruuzmCyjdV1KSqX/Eux04HfCWYmdxN3SQ/qqw=
 *         </G>
 *         <Y>
 *           pA5NnZvcd574WRXuOA7ZfC/7Lqt4cB0MRLWtHubtJoVOao9ib5ry4rTk0r6ddnOv
 *           AIGKktutzK3ymvKleS3DOrwZQgJ+/BDWDW8kO9R66o6rdjiSobBi/0c2V1+dkqOg
 *           jFmKz395mvCOZGhC7fqAVhHat2EjGPMfgSZyABa7+1k=
 *         </Y>
 *       </DSAKeyValue>
 *     </KeyValue>
 *   </KeyInfo>
 * </Signature>
 *</Envelope>
 * </code></pre>
 */
public class GenEnveloped {

    //
    // Synopsis: java GenEnveloped [document] [output]
    //
    //    where "document" is the name of a file containing the XML document
    //    to be signed, and "output" is the name of the file to store the
    //    signed document. The 2nd argument is optional - if not specified,
    //    standard output will be used.
    //
    public static void signXML(String inputFile, String outputFile, PrivateKey key, Certificate[] chain) throws Exception {
        // Create a DOM XMLSignatureFactory that will be used to generate the 
	// enveloped signature
        String providerName = System.getProperty
            ("jsr105Provider", "org.jcp.xml.dsig.internal.dom.XMLDSigRI");
	XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM",
	     (Provider) Class.forName(providerName).newInstance());

        // Create a Reference to the enveloped document (in this case we are
	// signing the whole document, so a URI of "" signifies that) and
	// also specify the SHA1 digest algorithm and the ENVELOPED Transform.
	Reference ref = fac.newReference
	    ("", fac.newDigestMethod(DigestMethod.SHA1, null),
             Collections.singletonList
	      (fac.newTransform
		(Transform.ENVELOPED, (TransformParameterSpec) null)), 
	     null, null);

	// Create the SignedInfo
	SignedInfo si = fac.newSignedInfo
	    (fac.newCanonicalizationMethod
	     (CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS, 
	      (C14NMethodParameterSpec) null), 
	     fac.newSignatureMethod(SignatureMethod.DSA_SHA1, null),
	     Collections.singletonList(ref));

        // Create a DSA KeyPair
//        KeyPairGenerator kpg = KeyPairGenerator.getInstance("DSA");
//	kpg.initialize(512);
//        KeyPair kp = kpg.generateKeyPair();

        // Create a KeyValue containing the DSA PublicKey that was generated
	KeyInfoFactory kif = fac.getKeyInfoFactory();
        KeyValue kv = kif.newKeyValue(chain[0].getPublicKey());

	// Create a KeyInfo and add the KeyValue to it
        KeyInfo ki = kif.newKeyInfo(Collections.singletonList(kv));

	// Instantiate the document to be signed
	DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
	dbf.setNamespaceAware(true);
	Document doc = 
	    dbf.newDocumentBuilder().parse(new FileInputStream(inputFile));

        // Create a DOMSignContext and specify the DSA PrivateKey and
        // location of the resulting XMLSignature's parent element
	DOMSignContext dsc = new DOMSignContext
	    (key, doc.getDocumentElement());

	// Create the XMLSignature (but don't sign it yet)
	XMLSignature signature = fac.newXMLSignature(si, ki);

        // Marshal, generate (and sign) the enveloped signature
        signature.sign(dsc);

	// output the resulting document
	OutputStream os = new FileOutputStream(outputFile);
	
	TransformerFactory tf = TransformerFactory.newInstance();
	Transformer trans = tf.newTransformer();
	trans.transform(new DOMSource(doc), new StreamResult(os));
    }
    
    public static int sign_xml(String inputFile, String outputFile, PrivateKey key, Certificate[] chain) {
        try {
          XMLSignatureFactory factory = XMLSignatureFactory.getInstance("DOM");
          //reference inside SignedInfo tag
          Reference reference1 = factory.newReference("",
            factory.newDigestMethod(DigestMethod.SHA1, null),
            Collections.singletonList(
              factory.newTransform(Transform.ENVELOPED,
                (TransformParameterSpec)null)),
            null,null);    
          
          //SignedInfo tag
          SignedInfo signed_info = factory.newSignedInfo(
            factory.newCanonicalizationMethod(
              CanonicalizationMethod.INCLUSIVE,(C14NMethodParameterSpec)null),
            factory.newSignatureMethod(SignatureMethod.RSA_SHA1,null),
            Collections.singletonList(reference1)
          );    
          
          //KeyInfo tag
          X509Certificate x509_cert = (X509Certificate) chain[0];
          KeyInfoFactory key_inf_fac = factory.getKeyInfoFactory();
          java.util.List<Object> x509_content = new ArrayList<Object>();
          x509_content.add(x509_cert.getSubjectX500Principal().getName());
          x509_content.add(x509_cert);
          X509Data x509_data = key_inf_fac.newX509Data(x509_content);
          KeyInfo key_info = 
            key_inf_fac.newKeyInfo(Collections.singletonList(x509_data));
          
          //get xml source and create signing context
          DocumentBuilderFactory doc_fac = DocumentBuilderFactory.newInstance();
          doc_fac.setNamespaceAware(true);
//          Document doc = doc_fac.newDocumentBuilder().parse(
//            new ByteArrayInputStream(file.read_all_bytes(file_path)));
          
          Document doc = 
        	  doc_fac.newDocumentBuilder().parse(new FileInputStream(inputFile));
          
          DOMSignContext context = new DOMSignContext(
            key,doc.getDocumentElement());
          
          //sign an xml signature
          XMLSignature xml_signature=factory.newXMLSignature(signed_info,key_info);
          signed_info.toString();
          xml_signature.sign(context);
          
          
          //save back to xml file
          OutputStream out_stream = new FileOutputStream(outputFile);
          TransformerFactory trans_fac = TransformerFactory.newInstance();
          Transformer trans = trans_fac.newTransformer();
          trans.transform(new DOMSource(doc),new StreamResult(out_stream));
          out_stream.close();
          return ErrorCodes.SIGN_XML.SUCCESS;
        }   
        catch (Exception e) {
          System.out.println("[Exception]" + e.getMessage());
          return ErrorCodes.SIGN_XML.XMLSignatureException;
        }
      }  
      
}
